package scu.maqiang.csparsej;

import scu.maqiang.csparsej.tdouble.Dcs_common;
import scu.maqiang.csparsej.tdouble.Dcs_util;
import scu.maqiang.numeric.ParamCheck;

import java.util.Arrays;

public class JSparse {
    /**
     * 元素最大个数
     */
    int nzmax;
    /**
     * 矩阵行数
     */
    int m;
    /**
     * 矩阵列数
     */
    int n;

    /**
     * 矩阵列指示向量(大小为n+1), 或者列指标(大小为nzmax)
     */
    int[] p;

    /**
     * 行指标, 大小为nzmax
     */
    int[] i;

    /**
     * 元素值, 大小为nzmax
     */
    double[] x;

    /**
     * 在三元素矩阵中为元素的个数, 在列压缩矩阵中为-1
     */
    int nz;

    /**
     * 构造函数, 给出指定大小稀疏矩阵, 指定可以包含的最大的元素个数
     * @param m 矩阵行数
     * @param n 矩阵列数
     * @param nzmax 可容纳的最大元素个数
     * @param values 如果为true, 则分配数值元素, 如果为false, 则不分配数值元素
     * @param triplet 为true, 则为三元组形式, 否则为列压缩形式
     */
    public JSparse(int m, int n, int nzmax, boolean values, boolean triplet) {
        this.m = m;
        this.n = n;
        this.nzmax = nzmax;
        this.nz = triplet?0:-1;
        this.p = new int[triplet ? nzmax : n + 1];
        this.i = new int[nzmax];
        this.x = values ? new double[nzmax] : null;
    }

    /**
     * 设置矩阵可包含的最大元素个数
     * @param nzmax
     */
    public boolean setNzmax(int nzmax) {
        if (nzmax <= 0)
            nzmax = isCompressedForm() ? p[n] : nz;
        int[] Ainew = new int[nzmax];
        int length = Math.min(nzmax, i.length);
        System.arraycopy(i, 0, Ainew, 0, length);
        i = Ainew;
        if (isTripletForm()) {
            int[] Apnew = new int[nzmax];
            length = Math.min(nzmax, p.length);
            System.arraycopy(p, 0, Apnew, 0, length);
            p = Apnew;
        }
        if (x != null) {
            double[] Axnew = new double[nzmax];
            length = Math.min(nzmax, x.length);
            System.arraycopy(x, 0, Axnew, 0, length);
            x = Axnew;
        }
        this.nzmax = nzmax;
        return true;
    }

    /**
     * 在三元组形式下稀疏矩阵元素, 如果矩阵为列压缩形式则无效, 此方法中要求所设行列指标值大于0,
     * 但不检查矩阵行列指标的上限, 如果行列指标的值超过当前矩阵的行数或者列数, 则矩阵的行数或者列数更新为此时的给出的行指标或列指标
     * @param i 行指标, 如果行指标大于行数, 则将行数设为当前行指标
     * @param j 列指标, 如果列指标大于列数, 则将行数设为当前列指标
     * @param val 第i行第j列元素值
     * @return 设置成功返回true; 若矩阵不为三元素形式并且行列指标均小于0，则返回false
     */
    public boolean setEntry(int i, int j, double val) {
        if(!isTripletForm() || i < 0 || j < 0) {
            return false;
        }
        if(nz >= nzmax && !setNzmax(2 * nzmax)) {
            return false;
        }
        if (x != null) {
            x[nz] = val;
        }
        this.i[nz] = i;
        this.p[nz++] = j;
        m = Math.max(m, i + 1);
        n = Math.max(n, j + 1);
        return true;
    }

    /**
     *
     * @param i
     * @param j
     * @param val
     * @return
     */
    public boolean setEntry2(int i, int j, double val) {
        if(!isTripletForm()) {
            return false;
        }
        ParamCheck.checkBounds(i, 0, m);
        ParamCheck.checkBounds(j, 0, n);
        if(nz >= nzmax && !setNzmax(2 * nzmax)) {
            return false;
        }
        if (x != null) {
            x[nz] = val;
        }
        this.i[nz] = i;
        this.p[nz++] = j;
        return true;
    }


    public boolean isCompressedForm() {
        return this.nz == -1;
    }

    public boolean isTripletForm() {
        return this.nz >= 0;
    }

    public void printMatrix(boolean brief) {
        if (nz < 0) {
            System.out.printf("%d-by-%d, nzmax: %d nnz: %d\n", m, n, nzmax, p[n]);
            for (int j = 0; j < n; j++) {
                System.out.print(String.format("    col %d : locations %d to %d\n", j, p[j], p[j + 1] - 1));
                for (int k = p[j]; k < p[j + 1]; k++) {
                    System.out.printf("      %d : %g\n", i[k], x != null ? x[k] : 1);
                    if (brief && k > 20) {
                        System.out.print("  ...\n");
                        return;
                    }
                }
            }
        } else {
            System.out.printf("triplet: %d-by-%d, nzmax: %d nnz: %d\n", m, n, nzmax, nz);
            for (int k = 0; k < nz; k++) {
                System.out.print(String.format("    %d %d : %g\n", i[k], p[k], x != null ? x[k] : 1));
                if (brief && k > 20) {
                    System.out.print("  ...\n");
                    return;
                }
            }
        }
    }

    public JSparse compress() {
       if(!isTripletForm()) {
           return null;
       }
        JSparse C = new JSparse(m, n, nz, x != null, false); /* allocate result */
        int[] w = new int[n]; /* get workspace */
        for (int k = 0; k < nz; k++)
            w[p[k]]++; /* column counts */
//        System.out.println(nz);
//        System.out.println(Arrays.toString(p));
//        System.out.println(Arrays.toString(w));
        cumsum(C.p, w, n); /* column pointers */
//        System.out.println(Arrays.toString(C.p));
//        System.out.println("w = " + Arrays.toString(w));
        int pp;
        for (int k = 0; k < nz; k++) {
            C.i[pp = w[p[k]]++] = i[k]; /* A(i,j) is the pth entry in C */
            if (C.x != null)
                C.x[pp] = x[k];
            System.out.println("pp = " + pp);
        }
//        System.out.println(Arrays.toString(C.i));
//        System.out.println(Arrays.toString(C.x));
//        System.out.println("w = " + Arrays.toString(w));
        return C;
    }

    public boolean cumsum(int[] p, int[] c, int n) {
        int i, nz = 0;
        double nz2 = 0;
        if (p == null || c == null)
            return false; /* check inputs */
        for (i = 0; i < n; i++) {
            p[i] = nz;
            nz += c[i];
            nz2 += c[i]; /* also in double to avoid int overflow */
            c[i] = p[i]; /* also copy p[0..n-1] back into c[0..n-1]*/
        }
        p[n] = nz;
        return (int) nz2 > 0; /* return sum (c [0..n-1]) */
    }

    public JSparse transpose(boolean values) {
        if (!isCompressedForm())
            return (null); /* check inputs */
        JSparse C = new JSparse(n, m, p[n], values && (x != null), false); /* allocate result */
        int[] w = new int[m]; /* get workspace */
        for (int pp = 0; pp < p[n]; pp++)
            w[i[pp]]++; /* row counts */
        cumsum(C.p, w, m); /* row pointers */
        int q = 0;
        for (int j = 0; j < n; j++) {
            for (int pp = p[j]; pp < p[j + 1]; pp++) {
                C.i[q = w[i[pp]]++] = j; /* place A(i,j) as entry C(j,i) */
                if (C.x != null)
                    C.x[q] = x[pp];
            }
        }
        return C;
    }

    public void duplicate() {
        if (!isCompressedForm())
            throw new IllegalArgumentException("Not CSC format!");
        /* check inputs */
        int[] w = new int[m]; /* get workspace */
        Arrays.fill(w, -1);
        int nnz = 0;
        for (int j = 0; j < n; j++) {
            int q = nnz; /* column j will start at q */
            for (int pp = p[j]; pp < p[j + 1]; pp++) {
                int row = i[pp]; /* A(i,j) is nonzero */
                if (w[row] >= q) {
                    x[w[row]] += x[pp]; /* A(i,j) is a duplicate */
                } else {
                    w[row] = nnz; /* record where row i occurs */
                    i[nnz] = row; /* keep A(i,j) */
                    x[nnz++] = x[pp];
                }
            }
            p[j] = q; /* record start of column j */
        }
        p[n] = nnz; /* finalize A */
        setNzmax(0);
    }

    public void gaxpy(double[] x, double[] y) {
        if (!isCompressedForm() || x == null || y == null)
            throw new IllegalArgumentException("illegal argument!");
        int nx = x.length;
        ParamCheck.checkEqual(nx, n);
        int ny = y.length;
        ParamCheck.checkEqual(ny, m);
        for (int j = 0; j < n; j++) {
            for (int pp = p[j]; pp < p[j + 1]; pp++) {
                y[i[pp]] += x[pp] * x[j];
            }
        }
    }

    public double norm1() {
        if (!isCompressedForm() || x == null)
            return (-1); /* check inputs */
        double norm = 0;
        for (int j = 0; j < n; j++) {
            double s = 0;
            for (int pp = p[j]; pp < p[j + 1]; pp++)
                s += Math.abs(x[pp]);
            norm = Math.max(norm, s);
        }
        return (norm);
    }


    public static void main(String[] args) {
        JSparse mat = new JSparse(10, 10, 1000, true, true);
        mat.setEntry2(0, 0, 10);
        mat.setEntry2(0, 1, 5);
        mat.setEntry2(0, 2, 30);
        mat.setEntry2(0, 2, 10);
        mat.setEntry2(0, 5, 21);
        mat.setEntry2(0, 4, 11);
        mat.setEntry2(3, 2, 32);
        mat.setEntry2(2, 2, 22);
        mat.setEntry2(1, 1, 3);
        mat.printMatrix(false);
        JSparse compressMat = mat.compress();
        compressMat.printMatrix(false);
        compressMat.duplicate();
        compressMat.printMatrix(false);

        JSparse compressMatT = compressMat.transpose(true);
        compressMatT.printMatrix(false);

    }
}
